package 困难.双指针;

/**
 * 一个数组的 分数 定义为数组之和 乘以 数组的长度。
 * 比方说，[1, 2, 3, 4, 5] 的分数为 (1 + 2 + 3 + 4 + 5) * 5 = 75 。
 * 给你一个正整数数组 nums 和一个整数 k ，请你返回 nums 中分数 严格小于 k 的 非空整数子数组数目。
 * 子数组 是数组中的一个连续元素序列。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/count-subarrays-with-score-less-than-k
 */
public class 统计得分小于K的子数组数目_2302 {

    public static void main(String[] args) {

        System.out.println(countSubarrays(new int[]{9, 5, 3, 8, 4, 7, 2, 7, 4, 5, 4, 9, 1, 4, 8, 10, 8, 10, 4, 7}, 4));

    }

    /**
     * 双指针
     * 1，定义两个指针，left，right，left不变，right++，分数一定是
     * 递增的，当分数 > k 时，满足条件的子数组数目 + (right - left)
     * 2，窗口中分数的变化计算：
     * (curK / (right - left) + nums[right]) * (right - left + 1) --> right++
     * (curK / (right - left) - nums[left]) * (right - left - 1)  --> left++
     * 注意：left < nums.length 时，剩下的以 [left,nums.length-1] 开头的子数组都是满足要求的，等差数列求和
     */
    public static long countSubarrays(int[] nums, long k) {
        long cnt = 0;
        long curK = 0;
        int left = 0, right = 0;
        while (right < nums.length) {
            // 退出是，curK 计算的是 [left,right-1] 的分数
            while (right < nums.length && curK < k) {
                long nextK;
                if (left == right) {
                    nextK = nums[right];
                } else {
                    nextK = (curK / (right - left) + nums[right]) * (right - left + 1);
                }
                if (nextK < k) {
                    curK = nextK;
                    right++;
                } else {
                    break;
                }
            }
            cnt += right - left;

            if (right != left) curK = (curK / (right - left) - nums[left]) * (right - left - 1);
            left++;
        }

        // 注意：left < nums.length 时，剩下的以 [left,nums.length-1] 开头
        // 的子数组都是满足要求的，等差数列求和，需要用long
        long n = nums.length - left;
        cnt += n * n - n * (n - 1) / 2;
        return cnt;
    }


    /**
     * 优化
     * 分数可以由 curSum * 窗口中元素个数得到
     */
    public static long countSubarrays2(int[] nums, long k) {
        long cnt = 0, curSum = 0;
        int left = 0, right = 0;
        while (left < nums.length) {
            while (right < nums.length && (curSum + nums[right]) * (right - left + 1) < k) {
                curSum += nums[right];
                right++;
            }
            cnt += right - left;

            curSum -= nums[left];
            left++;
        }
        return cnt;
    }

}
